package scu.maqiang.heat;

import scu.maqiang.fes.BVPType;
import scu.maqiang.fes.FES1L21;
import scu.maqiang.mesh.Mesh1L2;
import scu.maqiang.mesh.ScalarFunc;
import scu.maqiang.numeric.Direct;
import scu.maqiang.numeric.NewIterSSolver;
import scu.maqiang.numeric.SRMatrix;

import java.util.Arrays;

public class Heat1DL21 {
	public Heat1DL21(Mesh1L2 mesh) {
		this.mesh = mesh;
		fs = new FES1L21(mesh);
		A = new SRMatrix(fs.GetNdof());
		M = new SRMatrix(fs.GetNdof());
		x = new double[fs.GetNdof()];
		RHS = new double[fs.GetNdof()];
		source = new double[fs.GetNdof()];
		flux1 = new double[fs.GetNdof()];
		flux2 = new double[fs.GetNdof()];
		gradient = new double[1][fs.GetNdof()];
	}

	public static void main(String[] args) {
		long begin = System.nanoTime();
		double theta = 0.5;
		double T = 4.0;
		double dt = T / 1000.0;
		Mesh1L2 mesh = new Mesh1L2();
		mesh.line(0, 1, 100);
		
		Heat1DL21 heat = new Heat1DL21(mesh);
		heat.GeneralHeatProblem(new BVPType[] {BVPType.COMMON, BVPType.CONSIST_MASS},
				1.0, 1.0, 0.0, 
				new double[] {theta, T, dt}, 
				new int[] {2}, 1.0, 
				null, 0.0, 
				null, 0.0, 0.0, 
				"MyHeat1DL21.dat");
		long end = System.nanoTime();
		System.out.println("Solve time: " + (end - begin) / 1.0e9 + " s");
	}
	
	public void GeneralHeatProblem(BVPType[] tp, double kxx, double capacity, double sourceValue, 
			double[] timeIntegration, 
			int[] DirichletBoundaryLabel, double fixedvalue,
			int[] NeumannBoundaryLabel, double fluxValue,  
			int[] RobinBoundaryLabel, double robinFluxValue, double convectValue, String reportFile) {
		double theta = timeIntegration[0];
		double T = timeIntegration[1];
		double dt = timeIntegration[2];
		fs.assembleStiff(new double[] {kxx}, tp[0], A);
		if (RobinBoundaryLabel != null) {
			fs.assembleBoundaryMass(new double[] {convectValue}, tp[1], A, RobinBoundaryLabel);
		}
		fs.assembleMass(new double[] {capacity}, tp[1], M);
		
		SRMatrix L = M.add(theta * dt, A);
		SRMatrix R = M.add((theta - 1) * dt, A); 
		if (Math.abs(sourceValue) > 0.0) {
			fs.assembleSource(new double[] {sourceValue}, tp[0], source);
		}
		if (NeumannBoundaryLabel != null) {
			fs.assembleFlux(new double[] {fluxValue}, tp[0], flux1, NeumannBoundaryLabel);
		}
		if (RobinBoundaryLabel != null) {
			fs.assembleFlux(new double[] {robinFluxValue}, tp[0], flux2, RobinBoundaryLabel);
		}
		
		boolean dirichletExist = DirichletBoundaryLabel != null;
		if (dirichletExist) {
			fs.applyBC_MBN(L, Direct.All, DirichletBoundaryLabel);
		}
		
		NewIterSSolver solver = new NewIterSSolver(L);
//		DirectSSolver solver = new DirectSSolver(L);
		double time = 0.0;
		mesh.toTecplot(reportFile, time, x);
		double[] Rx = new double[fs.GetNdof()];
		while(time < T) {
			time += dt;
			System.out.println("time = " + time);
			R.mul(x, Rx);
			Arrays.setAll(RHS, i -> Rx[i] + (flux1[i] + flux2[i] + source[i]) * dt);
			if (dirichletExist) {
				fs.applyBC_MBN(RHS, Direct.All, fixedvalue, DirichletBoundaryLabel);
			}
			solver.PCGSSOR(RHS, x, 1.5, 1);
			mesh.toTecplot(reportFile, time, x);
		}
		
	}
	
	public void GeneralHeatProblem(BVPType[] tp, ScalarFunc kxxFunc, double[] paramKxx, 
			ScalarFunc capacityFunc, double[] paramCapacity, 
			ScalarFunc sourceFunc, double[] paramSource, 
			double[] timeIntegration, 
			int[] DirichletBoundaryLabel, ScalarFunc fixedFunc, double[] paramFixed,
			int[] NeumannBoundaryLabel, ScalarFunc fluxFunc, double[] paramFlux, 
			int[] RobinBoundaryLabel, ScalarFunc robinFluxFunc, double[] paramRobinFlux, 
			ScalarFunc convectFunc, double[] paramConvectFunc, String reportFile) {
		double theta = timeIntegration[0];
		double T = timeIntegration[1];
		double dt = timeIntegration[2];
		
		fs.assembleStiff(new ScalarFunc[] {kxxFunc}, paramKxx, tp[0], A);
		if (RobinBoundaryLabel != null) {
			fs.assembleBoundaryMass(new ScalarFunc[] {convectFunc}, paramConvectFunc, tp[1], A, RobinBoundaryLabel);
		}
		fs.assembleMass(new ScalarFunc[] {capacityFunc}, paramCapacity, tp[1], M);
		
		
		SRMatrix L = M.add(theta * dt, A);
		SRMatrix R = M.add((theta - 1) * dt, A);
		
		paramSource[0] = 0.0;
		boolean sourceExist = sourceFunc != null;
		if (sourceExist) {
			fs.assembleSource(new ScalarFunc[] {sourceFunc}, paramSource, tp[0], source);
		}
		
		double[] source2 = new double[fs.GetNdof()];
		
		boolean NeumannExist = NeumannBoundaryLabel != null;
		if (NeumannExist) {
			paramFlux[0] = 0.0;
			fs.assembleFlux(new ScalarFunc[] {fluxFunc}, paramFlux, tp[0], flux1, NeumannBoundaryLabel);
		}
		
		double[] flux11 = new double[fs.GetNdof()];
		
		boolean RobinExist = RobinBoundaryLabel != null;
		if (RobinExist) {
			paramRobinFlux[0] = 0.0;
			fs.assembleFlux(new ScalarFunc[] {robinFluxFunc}, paramRobinFlux, tp[0], flux2, RobinBoundaryLabel);
		}
		
		double[] flux22 = new double[fs.GetNdof()];
		
		boolean dirichletExist = DirichletBoundaryLabel != null;
		if (dirichletExist) {
			fs.applyBC_MBN(L, Direct.All, DirichletBoundaryLabel);
		}
		
		NewIterSSolver solver = new NewIterSSolver(L);
		double time = 0.0;
		mesh.toTecplot(reportFile, time, x);
		double[] Rx = new double[fs.GetNdof()];
		while(time < T) {
			time += dt;
			System.out.println("time = " + time);
			R.mul(x, Rx);
			if(sourceExist) {
				paramSource[0] = time;
				fs.assembleSource(new ScalarFunc[] {sourceFunc}, paramSource, tp[0], source2);
			}
			
			if (NeumannExist) {
				paramFlux[0] = time;
				fs.assembleFlux(new ScalarFunc[] {fluxFunc}, paramFlux, tp[0], flux11, NeumannBoundaryLabel);
			}
			
			if (RobinExist) {
				paramRobinFlux[0] = time;
				fs.assembleFlux(new ScalarFunc[] {robinFluxFunc}, paramRobinFlux, tp[0], flux22, RobinBoundaryLabel);
			}
			
			for(int i = 0, n = fs.GetNdof(); i < n; i++) {
				RHS[i] = Rx[i] + (1 - theta) * (flux11[i] + flux22[i] + source2[i]) * dt + theta * (flux1[i] + flux2[i] + source[i]) * dt;
				flux1[i] = flux11[i];
				flux2[i] = flux22[i];
				source[i] = source2[i];
			}
			if (dirichletExist) {
				paramFixed[0] = time;
				fs.applyBC_MBN(RHS, Direct.All, new ScalarFunc[]{fixedFunc}, paramFixed, DirichletBoundaryLabel);
			}
			solver.PCGSSOR(RHS, x, 1.50, 1);
			mesh.toTecplot(reportFile, time, x);
		}
	}
	
	Mesh1L2 mesh;
	FES1L21 fs;
	SRMatrix A;
	SRMatrix M;
	double[] source;
	double[] flux1;
	double[] flux2;
	double[] RHS;
	double[] x;
	double[][] gradient;

}
